#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2025/9/28 11:52
# @Author  : pym
# @File    : train.py
# @Project : mediapipe_hands
# @Desc    : 训练脚本
# 根据已有的拍摄点集，训练模型的生成点集，期望得到对应的输入分布

import os
import torch
import numpy as np
from NIMBLELayer import NIMBLELayer
from utils import batch_to_tensor_device, save_textured_nimble, smooth_mesh
import pytorch3d
import pytorch3d.io
from pytorch3d.structures.meshes import Meshes
from pytorch3d.transforms import axis_angle_to_matrix

# =====================
# 配置
# =====================
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(device)
OUTPUT_FOLDER = "train_output"
os.makedirs(OUTPUT_FOLDER, exist_ok=True)

# =====================
# 日志配置
# =====================
log_file = os.path.join(OUTPUT_FOLDER, "training_log.txt")

def log_message(message):
    """同时写入文件并在需要时打印"""
    with open(log_file, "a") as f:
        f.write(message + "\n")

def align_to(pred, target):
    # 去中心化
    pred_centered = pred - pred.mean(dim=1, keepdim=True)
    target_centered = target - target.mean(dim=1, keepdim=True)
    # 尺度归一化
    scale = target_centered.norm(dim=(1,2), keepdim=True) / pred_centered.norm(dim=(1,2), keepdim=True)
    pred_aligned = pred_centered * scale
    return pred_aligned, target_centered

# =====================
# 加载 MediaPipe keypoints
# 假设单帧 keypoints 已保存为 .npy
# shape: (21, 3)
# keypoints_path = "keypoints_frame0.npy"
# keypoints = torch.tensor(np.load(keypoints_path), dtype=torch.float32, device=device).unsqueeze(0)  # shape: (1, 21, 3)

# OK
# keypoints_data = [[ 3.9254e+02,  3.7800e+02, -7.2520e-05],
#         [ 3.4313e+02,  3.3547e+02, -1.8412e+01],
#         [ 3.0895e+02,  2.8223e+02, -3.6607e+01],
#         [ 2.8362e+02,  2.4825e+02, -5.6847e+01],
#         [ 2.8474e+02,  2.1564e+02, -7.6886e+01],
#         [ 3.7816e+02,  2.1142e+02, -1.5435e+01],
#         [ 3.5744e+02,  1.6091e+02, -5.4877e+01],
#         [ 3.2625e+02,  1.7189e+02, -8.8553e+01],
#         [ 2.9960e+02,  1.9558e+02, -1.0719e+02],
#         [ 4.1415e+02,  2.0907e+02, -2.5735e+01],
#         [ 4.2101e+02,  1.4491e+02, -5.6284e+01],
#         [ 4.0191e+02,  1.1056e+02, -8.2455e+01],
#         [ 3.7513e+02,  8.6807e+01, -9.8646e+01],
#         [ 4.4553e+02,  2.2677e+02, -4.1228e+01],
#         [ 4.6654e+02,  1.6497e+02, -6.7175e+01],
#         [ 4.6228e+02,  1.2414e+02, -8.5128e+01],
#         [ 4.5273e+02,  8.7386e+01, -9.6273e+01],
#         [ 4.6707e+02,  2.6061e+02, -5.8498e+01],
#         [ 5.0207e+02,  2.1966e+02, -7.9439e+01],
#         [ 5.2264e+02,  1.8580e+02, -8.8509e+01],
#         [ 5.3993e+02,  1.5112e+02, -9.3708e+01]]
# zan
keypoints_data = [[ 3.0996e+02,  4.3781e+02, -7.0221e-04],
        [ 2.7781e+02,  3.7398e+02, -2.1873e+01],
        [ 2.8480e+02,  2.8893e+02, -2.6693e+01],
        [ 2.9620e+02,  2.2901e+02, -3.2823e+01],
        [ 2.9871e+02,  1.8440e+02, -3.3756e+01],
        [ 3.5976e+02,  2.6361e+02,  7.3824e+00],
        [ 3.6908e+02,  2.3126e+02, -2.7406e+01],
        [ 3.3493e+02,  2.7346e+02, -4.4732e+01],
        [ 3.2268e+02,  3.0088e+02, -5.0815e+01],
        [ 3.9404e+02,  2.9042e+02,  3.8178e+00],
        [ 3.9835e+02,  2.6667e+02, -3.1300e+01],
        [ 3.4600e+02,  3.1305e+02, -3.5438e+01],
        [ 3.3713e+02,  3.3215e+02, -2.8615e+01],
        [ 4.1956e+02,  3.2623e+02, -7.8724e+00],
        [ 4.2051e+02,  3.0857e+02, -4.7117e+01],
        [ 3.6393e+02,  3.5147e+02, -3.5376e+01],
        [ 3.5096e+02,  3.6981e+02, -1.6354e+01],
        [ 4.3968e+02,  3.6747e+02, -2.1865e+01],
        [ 4.3280e+02,  3.5388e+02, -4.5985e+01],
        [ 3.8764e+02,  3.7988e+02, -3.6610e+01],
        [ 3.7540e+02,  3.9184e+02, -2.2616e+01]]
keypoints = torch.tensor(keypoints_data, dtype=torch.float32, device=device).unsqueeze(0)  # shape: (1, 21, 3)

# =====================
# 加载 NIMBLE 模型
# =====================
pm_dict_name = r"assets/NIMBLE_DICT_9137.pkl"
tex_dict_name = r"assets/NIMBLE_TEX_DICT.pkl"
mano_vreg_path = r"assets/NIMBLE_MANO_VREG.pkl"

pm_dict = batch_to_tensor_device(np.load(pm_dict_name, allow_pickle=True), device)
tex_dict = batch_to_tensor_device(np.load(tex_dict_name, allow_pickle=True), device)
nimble_mano_vreg = None
if os.path.exists(mano_vreg_path):
    nimble_mano_vreg = batch_to_tensor_device(np.load(mano_vreg_path, allow_pickle=True), device)

nlayer = NIMBLELayer(pm_dict, tex_dict, device,
                     use_pose_pca=True, pose_ncomp=30, shape_ncomp=20,
                     nimble_mano_vreg=nimble_mano_vreg)
nlayer.eval()

# =====================
# 初始化优化参数
# =====================
pose_param = torch.zeros((1, 30), requires_grad=True, device=device)
shape_param = torch.zeros((1, 20), requires_grad=True, device=device)

# 全局旋转和平移
global_rot = torch.zeros((1, 3), requires_grad=True, device=device)   # 轴角表示
global_trans = torch.zeros((1, 3), requires_grad=True, device=device) # 平移向量

tex_param = torch.zeros((1, 10), device=device)  # 保持原 demo, 不优化
# 可选加入全局平移/旋转，如果需要对齐摄像头坐标系

# optimizer = torch.optim.Adam([pose_param, shape_param], lr=1e-2)
# 优化器要加入它们
optimizer = torch.optim.Adam(
    [pose_param, shape_param, global_rot, global_trans],
    lr=1e-2
)

# =====================
# 对应索引
# MediaPipe 21点 -> NIMBLE bone_joints
# demo里输出 bone_joints，通常前21个关节对应
selected_joints_idx = list(range(21))

# =====================
# 训练循环
# =====================
num_steps = 1500
lambda_pose = 1e-3
lambda_shape = 1e-3

print("start....")

for step in range(num_steps):
    skin_v, muscle_v, bone_v, bone_joints, tex_img = nlayer.forward(
        pose_param, shape_param, tex_param, handle_collision=True
    )

     # 应用全局旋转和平移
    rot_mat = axis_angle_to_matrix(global_rot)  # (1, 3, 3)
    bone_joints = torch.matmul(bone_joints, rot_mat.transpose(1, 2)) + global_trans.unsqueeze(1)
    skin_v = torch.matmul(skin_v, rot_mat.transpose(1, 2)) + global_trans.unsqueeze(1)

    # ======== v1 ==========
    # # bone_joints: shape (1, num_joints, 3)
    # pred_joints = bone_joints[:, selected_joints_idx, :]

    # # 关键点拟合损失
    # loss_kp = torch.mean((pred_joints - keypoints) ** 2)

    # 取出对应点
    # pred_joints = bone_joints[:, list(mediapipe_to_nimble.values()), :]
    # target_joints = keypoints[:, list(mediapipe_to_nimble.keys()), :]

    # loss_kp = torch.mean((pred_joints - target_joints) ** 2)
    # loss_joints = torch.nn.functional.mse_loss(nimble_selected, mediapipe_joints)


    # ======= v2 ===========
    # 建立映射
    # nimble_to_mediapipe = {
    #     0: 0,   # wrist  carpal -> mediapipe wrist

    #     # Thumb
    #     2: 1,   # pro1 -> thumb_mcp
    #     3: 2,   # dis1 -> thumb_pip
    #     4: 3,   # dis1_end -> thumb_tip (mediapipe有thumb_dip，但通常映射到dis1)
    #     # 注意：MediaPipe thumb有 mcp, ip, tip，NIMBLE只有 pro, dis, end → 可以粗略对应

    #     # Index
    #     6: 5,   # pro2 -> index_mcp
    #     7: 6,   # int2 -> index_pip
    #     8: 7,   # dis2 -> index_dip
    #     9: 8,   # dis2_end -> index_tip

    #     # Middle
    #     11: 9,   # pro3 -> middle_mcp
    #     12: 10,  # int3 -> middle_pip
    #     13: 11,  # dis3 -> middle_dip
    #     14: 12,  # dis3_end -> middle_tip

    #     # Ring
    #     16: 13,  # pro4 -> ring_mcp
    #     17: 14,  # int4 -> ring_pip
    #     18: 15,  # dis4 -> ring_dip
    #     19: 16,  # dis4_end -> ring_tip

    #     # Pinky
    #     21: 17,  # pro5 -> pinky_mcp
    #     22: 18,  # int5 -> pinky_pip
    #     23: 19,  # dis5 -> pinky_dip
    #     24: 20   # dis5_end -> pinky_tip
    # }

    # ========= v3 修改大拇指的映射 ===========
    nimble_to_mediapipe = {
        0: 0,   # wrist

        # Thumb
        1: 1,   # met1 -> thumb_cmc
        2: 2,   # pro1 -> thumb_mcp
        3: 3,   # dis1 -> thumb_ip
        4: 4,   # dis1_end -> thumb_tip

        # Index
        6: 5,   # pro2 -> index_mcp
        7: 6,   # int2 -> index_pip
        8: 7,   # dis2 -> index_dip
        9: 8,   # dis2_end -> index_tip

        # Middle
        11: 9,   # pro3 -> middle_mcp
        12: 10,  # int3 -> middle_pip
        13: 11,  # dis3 -> middle_dip
        14: 12,  # dis3_end -> middle_tip

        # Ring
        16: 13,  # pro4 -> ring_mcp
        17: 14,  # int4 -> ring_pip
        18: 15,  # dis4 -> ring_dip
        19: 16,  # dis4_end -> ring_tip

        # Pinky
        21: 17,  # pro5 -> pinky_mcp
        22: 18,  # int5 -> pinky_pip
        23: 19,  # dis5 -> pinky_dip
        24: 20   # dis5_end -> pinky_tip
    }

    # ====== 删除手掌上的四个关节 =========
    pred_joints = []
    target_joints = []
    for n_idx, m_idx in nimble_to_mediapipe.items():
        pred_joints.append(bone_joints[:, n_idx, :])
        target_joints.append(keypoints[:, m_idx, :])
    pred_joints = torch.stack(pred_joints, dim=1)
    target_joints = torch.stack(target_joints, dim=1)

    # loss_kp = torch.mean((pred_joints - target_joints) ** 2)

    # ====== 去中心化 ========
    pred_joints_aligned, target_joints_aligned = align_to(pred_joints, target_joints)
    loss_kp = torch.mean((pred_joints_aligned - target_joints_aligned) ** 2)


    # 正则项
    loss_reg = lambda_pose * torch.mean(pose_param ** 2) + lambda_shape * torch.mean(shape_param ** 2)
    loss = loss_kp + loss_reg

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    if step % 50 == 0:
        print(f"Step {step}/{num_steps}, Loss: {loss.item():.6f}")
        log_message(f"Step {step}/{num_steps}, Loss: {loss.item():.6f}")

print("finish....")

# =====================
# 输出拟合结果
# =====================
skin_v, muscle_v, bone_v, bone_joints, tex_img = nlayer.forward(
    pose_param, shape_param, tex_param, handle_collision=True
)
skin_v_smooth = smooth_mesh(Meshes(skin_v, nlayer.skin_f.repeat(1, 1, 1))).verts_padded().detach().cpu().numpy()
bone_joints_np = bone_joints.detach().cpu().numpy()

# 保存骨骼关节
np.savetxt(os.path.join(OUTPUT_FOLDER, "fitted_bone_joints.xyz"), bone_joints_np[0])
# 保存 MANO 对应表面顶点
skin_mano_v = nlayer.nimble_to_mano(skin_v, is_surface=True)
np.savetxt(os.path.join(OUTPUT_FOLDER, "fitted_skin_manov.xyz"), skin_mano_v[0].detach().cpu().numpy())

# 保存网格 OBJ
skin_p3dmesh = Meshes(torch.tensor(skin_v_smooth, dtype=torch.float32, device=device),
                      nlayer.skin_f.repeat(1, 1, 1))
muscle_p3dmesh = Meshes(muscle_v, nlayer.muscle_f.repeat(1, 1, 1))
bone_p3dmesh = Meshes(bone_v, nlayer.bone_f.repeat(1, 1, 1))

pytorch3d.io.IO().save_mesh(skin_p3dmesh[0], os.path.join(OUTPUT_FOLDER, "fitted_skin.obj"))
pytorch3d.io.IO().save_mesh(muscle_p3dmesh[0], os.path.join(OUTPUT_FOLDER, "fitted_muscle.obj"))
pytorch3d.io.IO().save_mesh(bone_p3dmesh[0], os.path.join(OUTPUT_FOLDER, "fitted_bone.obj"))

# 保存纹理
tex_img_np = tex_img.detach().cpu().numpy()
save_textured_nimble(os.path.join(OUTPUT_FOLDER, "fitted_skin_textured.obj"), skin_v_smooth[0], tex_img_np[0])

print("NIMBLE拟合完成，结果已保存到:", OUTPUT_FOLDER)

